from langgraph.constants import START
from langgraph.graph import StateGraph, END
from typing import Dict, Any, TypedDict
from .nodes import RAGNodes


class RAGState(TypedDict):
    """RAG状态定义"""
    query: str
    retrieved_documents: list
    retrieval_scores: list
    answer: str


class RAGGraph:
    """RAG工作流图"""

    def __init__(self, llm, vector_store):
        self.nodes = RAGNodes(llm, vector_store)
        self.graph = self._build_graph()

    def _route_query(self, state: Dict[str, Any]) -> str:
        """路由查询的条件函数"""
        query = state.get("query", "")
        if len(query.strip()) == 0:
            return "direct_answer"
        return "retrieve"

    def _build_graph(self) -> StateGraph:
        """构建RAG工作流图"""
        workflow = StateGraph(RAGState)

        # 添加节点
        workflow.add_node("retrieve", self.nodes.retrieve_documents)
        workflow.add_node("generate", self.nodes.generate_answer)
        workflow.add_node("direct", self.nodes.direct_answer)

        # # 设置入口点
        # workflow.set_entry_point("check_query")

        # 添加条件边
        workflow.add_conditional_edges(
            START,
            self._route_query,
            {
                "retrieve": "retrieve",
                "direct_answer": "direct"
            }
        )

        # 添加边
        workflow.add_edge("retrieve", "generate")
        workflow.add_edge("generate", END)
        workflow.add_edge("direct", END)

        return workflow.compile()

    def invoke(self, query: str) -> Dict[str, Any]:
        """执行RAG查询"""
        initial_state = {
            "query": query,
            "retrieved_documents": [],
            "retrieval_scores": [],
            "answer": ""
        }
        # 使用这种的话需要其中所有节点都要按照字典的形式输出
        result = self.graph.invoke(initial_state)
        # return result
        # 如果答案是生成器，则直接返回
        if callable(result.get("answer")):
            return result
        else:
            return result